from mbsh.core.images import read_img_from_base64
from mbsh.core.models import get_img_dir
from tests import show_img

__author__ = 'gpu'

import unittest
import os

from mbsh import create_app, g_temp_dir
from mbsh.core.plot import Plot
import cv2


class PlotTestCase(unittest.TestCase):
    def setUp(self):
        app = create_app(os.getenv('FLASK_CONFIG') or 'default')
        app.app_context().push()
        self.app = app

    def tearDown(self):
        print('tear down')

    def test_plot(self):
        Plot.show_matrix([1, 0], [0, 1], 2, g_temp_dir)

    def test_bbox(self):
        cam = cv2.imread(r'E:\projects\znyx\znyx-model\tests\res\heatmap2.png', 0)
        img = cv2.imread(r'E:\projects\znyx\znyx-model\tests\res\heatmap2.png')
        # img=cv2.imread(r'E:\projects\znyx\znyx-model\tests\res\transparent.png')

        overlayBox, file = Plot.plot_cam(img, cam, get_img_dir())
        img = read_img_from_base64(request_data={"imgBase64": overlayBox})
        show_img(img)

    def test_cal_bbox(self):
        cam = cv2.imread(r'E:\projects\znyx\znyx-model\tests\res\heatmap2.png', 0)
        img = cv2.imread(r'E:\projects\znyx\znyx-model\tests\res\heatmap2.png')
        # img=cv2.imread(r'E:\projects\znyx\znyx-model\tests\res\transparent.png')

        rect = Plot.cal_cam_bbox(cam)
        self.assertEqual([56, 36, 168, 148], rect)

        img_path = Plot.plot_cam_bbox(img, [rect], 0.97, g_temp_dir)
        show_img(cv2.imread(img_path))

        img_path = Plot.plot_cam_bbox(img, [rect], "", g_temp_dir, color=(255, 0, 0))
        show_img(cv2.imread(img_path))


if __name__ == '__main__':
    unittest.main()
